//      vom.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef VARIABLE_ORDER_MODEL_H
#define VARIABLE_ORDER_MODEL_H

#include "stl-helper.h"
#include "dna_processing.h"
#include "suffix_array.h"

class VOMNode 
{
  friend class VOM;
 public:
  VOMNode() 
    : c_('X'), parent_(NULL) {}
  VOMNode(char c, VOMNode* parent)
    : c_(c), parent_(parent) {}

  ~VOMNode() {
    BOOST_FOREACH(VOMNode* child, child_node_)
      delete child;
  }

  void add_child(VOMNode* child) {
    child_char_.pb(child->c_);
    child_node_.pb(child);
  }

  VOMNode* child(char c) {
    for (uint i = 0; i < child_char_.size(); i++)
      if (child_char_[i] == c) return child_node_[i];
    return NULL;
  }

  std::string str() {
    std::string s;
    VOMNode* node = this;
    while (node->parent_ != NULL) {
      s = node->c_ + s;
      node = node->parent_;
    }
    return s;
  }

  char c_;
  std::vector<char> child_char_;
  VOMNode* parent_;
  std::vector<VOMNode*> child_node_;
  int count_;

  // transition probabilities
  vd_t prob_;
};

class VOM 
{
 public:
  VOM();
  ~VOM() { delete root_; }

  void build_big_tree(const Sequences& seqs, const SequencesSA& seqsa) { 
    root_->count_ = 0;
    BOOST_FOREACH(const std::string& s, seqs)
      root_->count_ += s.size();
    build_big_tree(root_, seqsa); 
  }

  void build_fix_tree(const Sequences& seqs, const SequencesSA& seqsa, int order)
  {
    root_->count_ = 0;
    BOOST_FOREACH(const std::string& s, seqs)
      root_->count_ += s.size();
    build_fix_tree(root_, seqsa, order+1);    
  }

  void mle_estimate() { 
    mle_estimate(root_, root_->child('A'), root_->child('C'), root_->child('G'), root_->child('T'));
  }

  std::ostream& print(std::ostream& o, bool isFile = false) const;
  std::ostream& print(std::ostream& o, VOMNode* node, int depth) const;

  void load(const std::string& modelFile);

  double log_likelihood(const std::string& seq, int max_order = -1) const;
 protected:
  void mle_estimate(VOMNode* node, VOMNode* na, VOMNode* nc, VOMNode* ng, VOMNode* nt);
  void build_big_tree(VOMNode* node, const SequencesSA& seqsa);
  void build_fix_tree(VOMNode* node, const SequencesSA& seqsa, int remain_depth);
  const vd_t& prob_vec(const std::string& seq, int i, int max_order) const;

  VOMNode* root_;

  std::string tmp_s;
};

#endif
